"""逻辑回归量化模型"""
import pandas as pd
import numpy as np
from sklearn.linear_model import LogisticRegression
from typing import Dict, Optional
from .base_model import BaseQuantModel


class LogisticRegressionModel(BaseQuantModel):
    """逻辑回归模型（分类：上涨/下跌）"""
    
    def __init__(self, **kwargs):
        """
        初始化逻辑回归模型
        
        Args:
            **kwargs: LogisticRegression参数
                常用参数：
                - C: 正则化强度，默认1.0
                - max_iter: 最大迭代次数，默认1000
                - solver: 求解器，默认'lbfgs'
        """
        super().__init__(model_name="LogisticRegression", task_type="classification")
        self.model_params = kwargs if kwargs else {'C': 1.0, 'max_iter': 1000, 'solver': 'lbfgs'}
    
    def _build_model(self, **kwargs):
        """构建逻辑回归模型"""
        params = {**self.model_params, **kwargs}
        self.model = LogisticRegression(**params, random_state=42)
    
    def _train_model(
        self,
        X_train: pd.DataFrame,
        y_train: pd.Series,
        X_val: Optional[pd.DataFrame] = None,
        y_val: Optional[pd.Series] = None,
        **kwargs
    ) -> Dict:
        """训练逻辑回归模型"""
        self.model.fit(X_train, y_train)
        
        # 计算训练集准确率
        train_score = self.model.score(X_train, y_train)
        
        result = {
            'train_score': train_score
        }
        
        # 如果有验证集，计算验证集准确率
        if X_val is not None and y_val is not None:
            val_score = self.model.score(X_val, y_val)
            result['val_score'] = val_score
        
        return result
    
    def _predict_proba(self, X: pd.DataFrame) -> np.ndarray:
        """预测概率"""
        return self.model.predict_proba(X)

